package com.element.oauth2.server.global.provider;

import com.element.oauth2.server.global.generator.OAuthAuthorizationCodeGenerator;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.security.authentication.AnonymousAuthenticationToken;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.crypto.keygen.Base64StringKeyGenerator;
import org.springframework.security.crypto.keygen.StringKeyGenerator;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.OAuth2Error;
import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
import org.springframework.security.oauth2.server.authorization.*;
import org.springframework.security.oauth2.server.authorization.authentication.*;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.security.Principal;
import java.util.Base64;
import java.util.Set;
import java.util.function.Consumer;

public class OAuthCodeRequestAuthorizationProvider implements AuthenticationProvider {

    private static final StringKeyGenerator DEFAULT_STATE_GENERATOR = new Base64StringKeyGenerator(Base64.getUrlEncoder());
    private final Log logger = LogFactory.getLog(this.getClass());
    private final RegisteredClientRepository registeredClientRepository;
    private final OAuth2AuthorizationService authorizationService;
    private final OAuth2AuthorizationConsentService authorizationConsentService;
    private final OAuth2TokenGenerator<OAuth2AuthorizationCode> authorizationCodeGenerator = new OAuthAuthorizationCodeGenerator();
    private final Consumer<OAuth2AuthorizationCodeRequestAuthenticationContext> authenticationValidator = new OAuth2AuthorizationCodeRequestAuthenticationValidator();

    public OAuthCodeRequestAuthorizationProvider(RegisteredClientRepository registeredClientRepository,
                                                 OAuth2AuthorizationService authorizationService,
                                                 OAuth2AuthorizationConsentService authorizationConsentService) {
        Assert.notNull(registeredClientRepository, "registeredClientRepository cannot be null");
        Assert.notNull(authorizationService, "authorizationService cannot be null");
        Assert.notNull(authorizationConsentService, "authorizationConsentService cannot be null");
        this.registeredClientRepository = registeredClientRepository;
        this.authorizationService = authorizationService;
        this.authorizationConsentService = authorizationConsentService;
    }

    public Authentication authenticate(Authentication authentication) throws AuthenticationException {
        OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthentication = (OAuth2AuthorizationCodeRequestAuthenticationToken) authentication;
        RegisteredClient registeredClient = this.registeredClientRepository.findByClientId(authorizationCodeRequestAuthentication.getClientId());
        if (registeredClient == null) {
            throwError("invalid_request", authorizationCodeRequestAuthentication, null);
        }

        if (this.logger.isTraceEnabled()) {
            this.logger.trace("Retrieved registered client");
        }

        OAuth2AuthorizationCodeRequestAuthenticationContext authenticationContext = OAuth2AuthorizationCodeRequestAuthenticationContext.with(authorizationCodeRequestAuthentication).registeredClient(registeredClient).build();
        this.authenticationValidator.accept(authenticationContext);
        assert registeredClient != null;
        if (!registeredClient.getAuthorizationGrantTypes().contains(AuthorizationGrantType.AUTHORIZATION_CODE)) {
            throwError("unauthorized_client", authorizationCodeRequestAuthentication, registeredClient);
        }

        String codeChallenge = (String) authorizationCodeRequestAuthentication.getAdditionalParameters().get("code_challenge");
        if (StringUtils.hasText(codeChallenge)) {
            String codeChallengeMethod = (String) authorizationCodeRequestAuthentication.getAdditionalParameters().get("code_challenge_method");
            if (!StringUtils.hasText(codeChallengeMethod) || !"S256".equals(codeChallengeMethod)) {
                throwError("invalid_request", "code_challenge_method", "https://datatracker.ietf.org/doc/html/rfc7636#section-4.4.1", authorizationCodeRequestAuthentication, registeredClient, (OAuth2AuthorizationRequest) null);
            }
        } else if (registeredClient.getClientSettings().isRequireProofKey()) {
            throwError("invalid_request", "code_challenge", "https://datatracker.ietf.org/doc/html/rfc7636#section-4.4.1", authorizationCodeRequestAuthentication, registeredClient, (OAuth2AuthorizationRequest) null);
        }

        if (this.logger.isTraceEnabled()) {
            this.logger.trace("Validated authorization code request parameters");
        }

        Authentication principal = (Authentication) authorizationCodeRequestAuthentication.getPrincipal();
        if (!isPrincipalAuthenticated(principal)) {
            if (this.logger.isTraceEnabled()) {
                this.logger.trace("Did not authenticate authorization code request since principal not authenticated");
            }
            return authorizationCodeRequestAuthentication;
        } else {
            OAuth2AuthorizationRequest authorizationRequest = OAuth2AuthorizationRequest.authorizationCode().authorizationUri(authorizationCodeRequestAuthentication.getAuthorizationUri()).clientId(registeredClient.getClientId()).redirectUri(authorizationCodeRequestAuthentication.getRedirectUri()).scopes(authorizationCodeRequestAuthentication.getScopes()).state(authorizationCodeRequestAuthentication.getState()).additionalParameters(authorizationCodeRequestAuthentication.getAdditionalParameters()).build();
            OAuth2AuthorizationConsent currentAuthorizationConsent = this.authorizationConsentService.findById(registeredClient.getId(), principal.getName());
            if (requireAuthorizationConsent(registeredClient, authorizationRequest, currentAuthorizationConsent)) {
                String state = DEFAULT_STATE_GENERATOR.generateKey();
                OAuth2Authorization authorization = authorizationBuilder(registeredClient, principal, authorizationRequest).attribute("state", state).build();
                if (this.logger.isTraceEnabled()) {
                    this.logger.trace("Generated authorization consent state");
                }

                this.authorizationService.save(authorization);
                Set<String> currentAuthorizedScopes = currentAuthorizationConsent != null ? currentAuthorizationConsent.getScopes() : null;
                if (this.logger.isTraceEnabled()) {
                    this.logger.trace("Saved authorization");
                }

                return new OAuth2AuthorizationConsentAuthenticationToken(authorizationRequest.getAuthorizationUri(), registeredClient.getClientId(), principal, state, currentAuthorizedScopes, null);
            } else {
                OAuth2TokenContext tokenContext = createAuthorizationCodeTokenContext(authorizationCodeRequestAuthentication, registeredClient, authorizationRequest.getScopes());
                OAuth2AuthorizationCode authorizationCode = (OAuth2AuthorizationCode) this.authorizationCodeGenerator.generate(tokenContext);
                if (authorizationCode == null) {
                    OAuth2Error error = new OAuth2Error("server_error", "The token generator failed to generate the authorization code.", "https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.2.1");
                    throw new OAuth2AuthorizationCodeRequestAuthenticationException(error, (OAuth2AuthorizationCodeRequestAuthenticationToken) null);
                } else {
                    if (this.logger.isTraceEnabled()) {
                        this.logger.trace("Generated authorization code");
                    }

                    OAuth2Authorization authorization = authorizationBuilder(registeredClient, principal, authorizationRequest).authorizedScopes(authorizationRequest.getScopes()).token(authorizationCode).build();
                    this.authorizationService.save(authorization);
                    if (this.logger.isTraceEnabled()) {
                        this.logger.trace("Saved authorization");
                    }

                    String redirectUri = authorizationRequest.getRedirectUri();
                    if (!StringUtils.hasText(redirectUri)) {
                        redirectUri = registeredClient.getRedirectUris().iterator().next();
                    }

                    if (this.logger.isTraceEnabled()) {
                        this.logger.trace("Authenticated authorization code request");
                    }

                    return new OAuth2AuthorizationCodeRequestAuthenticationToken(authorizationRequest.getAuthorizationUri(), registeredClient.getClientId(), principal, authorizationCode, redirectUri, authorizationRequest.getState(), authorizationRequest.getScopes());
                }
            }
        }
    }

    public boolean supports(Class<?> authentication) {
        return OAuth2AuthorizationCodeRequestAuthenticationToken.class.isAssignableFrom(authentication);
    }

    private static OAuth2Authorization.Builder authorizationBuilder(RegisteredClient registeredClient, Authentication principal, OAuth2AuthorizationRequest authorizationRequest) {
        return OAuth2Authorization.withRegisteredClient(registeredClient).principalName(principal.getName()).authorizationGrantType(AuthorizationGrantType.AUTHORIZATION_CODE).attribute(Principal.class.getName(), principal).attribute(OAuth2AuthorizationRequest.class.getName(), authorizationRequest);
    }

    private static OAuth2TokenContext createAuthorizationCodeTokenContext(OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthentication, RegisteredClient registeredClient, Set<String> authorizedScopes) {
        DefaultOAuth2TokenContext.Builder tokenContextBuilder = DefaultOAuth2TokenContext.builder()
                .registeredClient(registeredClient)
                .principal((Authentication) authorizationCodeRequestAuthentication.getPrincipal())
                .authorizationServerContext(AuthorizationServerContextHolder.getContext())
                .tokenType(new OAuth2TokenType("code")).authorizedScopes(authorizedScopes)
                .authorizationGrantType(AuthorizationGrantType.AUTHORIZATION_CODE)
                .authorizationGrant(authorizationCodeRequestAuthentication);
        return tokenContextBuilder.build();
    }

    private static boolean requireAuthorizationConsent(RegisteredClient registeredClient, OAuth2AuthorizationRequest authorizationRequest, OAuth2AuthorizationConsent authorizationConsent) {
        if (!registeredClient.getClientSettings().isRequireAuthorizationConsent()) {
            return false;
        } else if (authorizationRequest.getScopes().contains("openid") && authorizationRequest.getScopes().size() == 1) {
            return false;
        } else {
            return authorizationConsent == null || !authorizationConsent.getScopes().containsAll(authorizationRequest.getScopes());
        }
    }

    private static boolean isPrincipalAuthenticated(Authentication principal) {
        return principal != null && !AnonymousAuthenticationToken.class.isAssignableFrom(principal.getClass()) && principal.isAuthenticated();
    }

    private static void throwError(String errorCode, OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthentication, RegisteredClient registeredClient) {
        throwError(errorCode, "client_id", "https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.2.1", authorizationCodeRequestAuthentication, registeredClient, null);
    }

    private static void throwError(String errorCode, String parameterName, String errorUri, OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthentication, RegisteredClient registeredClient, OAuth2AuthorizationRequest authorizationRequest) {
        OAuth2Error error = new OAuth2Error(errorCode, "OAuth 2.0 Parameter: " + parameterName, errorUri);
        throwError(error, parameterName, authorizationCodeRequestAuthentication, registeredClient, authorizationRequest);
    }

    private static void throwError(OAuth2Error error, String parameterName, OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthentication, RegisteredClient registeredClient, OAuth2AuthorizationRequest authorizationRequest) {
        String redirectUri = resolveRedirectUri(authorizationCodeRequestAuthentication, authorizationRequest, registeredClient);
        if (error.getErrorCode().equals("invalid_request") && (parameterName.equals("client_id") || parameterName.equals("state"))) {
            redirectUri = null;
        }

        OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthenticationResult = new OAuth2AuthorizationCodeRequestAuthenticationToken(authorizationCodeRequestAuthentication.getAuthorizationUri(), authorizationCodeRequestAuthentication.getClientId(), (Authentication) authorizationCodeRequestAuthentication.getPrincipal(), redirectUri, authorizationCodeRequestAuthentication.getState(), authorizationCodeRequestAuthentication.getScopes(), authorizationCodeRequestAuthentication.getAdditionalParameters());
        throw new OAuth2AuthorizationCodeRequestAuthenticationException(error, authorizationCodeRequestAuthenticationResult);
    }

    private static String resolveRedirectUri(OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthentication, OAuth2AuthorizationRequest authorizationRequest, RegisteredClient registeredClient) {
        if (authorizationCodeRequestAuthentication != null && StringUtils.hasText(authorizationCodeRequestAuthentication.getRedirectUri())) {
            return authorizationCodeRequestAuthentication.getRedirectUri();
        } else if (authorizationRequest != null && StringUtils.hasText(authorizationRequest.getRedirectUri())) {
            return authorizationRequest.getRedirectUri();
        } else {
            return registeredClient != null ? registeredClient.getRedirectUris().iterator().next() : null;
        }
    }
}